


use actix_web::{get, middleware, post, web, App, Error, HttpResponse, HttpServer, HttpRequest};
use bytes::Buf;
use mime_guess::from_path;
use rust_embed::RustEmbed;
use std::borrow::Cow;
use std::collections::HashMap;
use std::io::Write;
use std::ops::Add;
use std::sync::{Arc, Mutex};
use actix::Actor;
use actix_web::web::Bytes;
use serde::{Deserialize, Serialize};
use diesel::prelude::*;
use diesel::r2d2::{self, ConnectionManager};
use uuid::Uuid;
use crate::config::ConfigAction;
use crate::ip_block::IpBlock;
use crate::kv::KVStore;
use crate::models::{Rule, YaraResult};
use crate::scheduler::Scheduler;
use crate::session::SessionStore;
use crate::virus_scheduler::VirusScheduler;
use chrono::prelude::*;
use std::process::{Command, Stdio};
use crate::falco_scheduler::FalcoScheduler;
use crate::suricata_scheduler::SuricataScheduler;
use actix_web_actors::ws;
use self::cpu_mem_ws::MyWebSocket;
use futures::stream::TryStreamExt; // Needed for the `.try_fold` method
use futures::StreamExt;
use actix_multipart::Multipart;
use aes::Aes128;
use aes::cipher::{
    BlockCipher, BlockEncrypt, BlockDecrypt, KeyInit,
    generic_array::GenericArray,
};
use service::*;
mod actions;
mod models;
mod schema;
mod kv;
mod config;
mod scheduler;
mod session;
mod redirect;
mod falco_scheduler;
mod suricata_scheduler;
mod cpu_mem_ws;
mod ip_block;
mod virus_scheduler;
mod service;
mod action;
mod schedule;
mod model;
mod response;
mod util;
mod plugin;

type DbPool = r2d2::Pool<ConnectionManager<SqliteConnection>>;

#[macro_use]
extern crate diesel;
#[macro_use]
extern crate lazy_static;

#[derive(Debug, Serialize, Deserialize)]
struct MyObj {
    name: String,
    number: i32,
}


#[derive(Debug, Serialize, Deserialize)]
struct MyObjResult {
    id: String,
    code: String,
    result: bool,
}

#[derive(RustEmbed)]
#[folder = "dist/"]
struct Asset;

fn handle_embedded_file(path: &str) -> HttpResponse {
    match Asset::get(path) {
        Some(content) => {
            let body: Bytes = match content.data {
                Cow::Borrowed(bytes) => bytes.into(),
                Cow::Owned(bytes) => bytes.into(),
            };
            HttpResponse::Ok().content_type(from_path(path).first_or_octet_stream().as_ref()).body(body)
        }
        None => HttpResponse::NotFound().body("404 Not Found"),
    }
}

async fn index() -> HttpResponse {
    handle_embedded_file("index.html")
}

async fn dist(path: web::Path<String>) -> HttpResponse {
    let url = &path.into_inner();
    if url.contains(".") {
        handle_embedded_file(url)
    } else {
        handle_embedded_file(url.clone().add("/index.html").as_str())
    }
}

/// Finds user by UID.
// #[get("/api/alert/{id}")]
// async fn get_user(
//     pool: web::Data<DbPool>,
//     id: web::Path<String>,
// ) -> Result<HttpResponse, Error> {
//     let id = id.into_inner();
//     let alert_id = id.clone();

//     // use web::block to offload blocking Diesel code without blocking server thread
//     let alert = web::block(move || {
//         let conn = pool.get()?;
//         actions::find_alert_by_id(alert_id, &conn)
//     })
//     .await?
//     .map_err(actix_web::error::ErrorInternalServerError)?;

//     if let Some(alert) = alert {
//         Ok(HttpResponse::Ok().json(alert))
//     } else {
//         let res = HttpResponse::NotFound().body(format!("No user found with uid: {}", id));
//         Ok(res)
//     }
// }

/// Inserts new user with name defined in form.
#[post("/api/alert/query/page")]
async fn find_alert_by_page(
    pool: web::Data<DbPool>,
    form: web::Json<models::AlertDTO>,
) -> Result<HttpResponse, Error> {
    // use web::block to offload blocking Diesel code without blocking server thread
    let alert_vec = web::block(move || {
        let conn = pool.get()?;
        actions::find_alert_by_page(form.current, form.page_size, form.source.as_ref(), &conn)
    })
    .await?
    .map_err(actix_web::error::ErrorInternalServerError)?;

    Ok(HttpResponse::Ok().json(models::ListResult {
        success: true,
        error_code: "".to_string(),
        error_message: "".to_string(),
        data: alert_vec.0,
        total: alert_vec.1,
    }))
}

/// Inserts new user with name defined in form.
#[post("/api/alert/solve/{id}")]
async fn solve_alert(
    pool: web::Data<DbPool>,
    ip_block: web::Data<ip_block::Iptables>,
    id: web::Path<String>,
) -> Result<HttpResponse, Error> {
    let id = id.into_inner();
    // use web::block to offload blocking Diesel code without blocking server thread
    let alert_vec = web::block(move || {
        let conn = pool.get()?;
        actions::solve_alert(id, ip_block, &conn)
    })
    .await?
    .map_err(actix_web::error::ErrorInternalServerError)?;

    Ok(HttpResponse::Ok().json(models::SysResult {
        success: true,
        error_code: "".to_string(),
        error_message: "".to_string(),
        data: Some(alert_vec),
    }))
}

/// Inserts new user with name defined in form.
#[post("/api/alert/query/kv")]
async fn alert_query_kv(
    rocksdb: web::Data<kv::RocksDB>
) -> Result<HttpResponse, Error> {
    let list = rocksdb.iterator();
    Ok(HttpResponse::Ok().json(list))
}


/// Inserts new user with name defined in form.
#[get("/api/virus/query/progress")]
async fn virus_query_progress(
    last_task: web::Data<Arc<Mutex<YaraResult>>>
) -> Result<HttpResponse, Error> {
    let yara_result = last_task.lock().unwrap();
    Ok(HttpResponse::Ok().json(models::SysResult {
        success: true,
        error_code: "".to_string(),
        error_message: "".to_string(),
        data: Some(yara_result.to_owned()),
    }))
}


/// Inserts new user with name defined in form.
#[get("/api/virus/start/scan")]
async fn virus_start_scan(
    rocksdb: web::Data<kv::RocksDB>
) -> Result<HttpResponse, Error> {
    match rocksdb.find("scanning") {
        Some(_) => {
            return Ok(HttpResponse::Ok().json(models::SysResult {
                success: true,
                error_code: "".to_string(),
                error_message: "病毒扫描任务正在进行".to_string(),
                data: Some(()),
            }));
        },
        None => {
            rocksdb.save("scanning", "");
            std::thread::spawn(move || {
                println!("task run");
                // Command::new(format!("{}/tailmoncli", "/root/projects/tailmon-edr/tailmoncli"))
                //     .current_dir("/root/projects/tailmon-edr/tailmoncli")
                //     .env("LD_LIBRARY_PATH", "$LD_LIBRARY_PATH:/usr/local/lib")
                //     .arg("-rule")
                //     .arg("/root/projects/tailmon-edr/tailmond/data/rules/current.yara")
                //     .arg("-threads")
                //     .arg("4")
                //     .arg("/root/")
                //     .status()
                //     .expect("tailmoncli command failed to start");
                // let mut child = Command::new(format!("{}/tailmoncli", "/root/projects/tailmon-edr/tailmoncli"))
                //     .current_dir("/root/projects/tailmon-edr/tailmoncli")
                //     .env("LD_LIBRARY_PATH", "$LD_LIBRARY_PATH:/usr/local/lib")
                //     .arg("-rule")
                //     .arg("/root/projects/tailmon-edr/tailmond/data/rules/current.yara")
                //     .arg("-threads")
                //     .arg("1")
                //     .arg("/root/")
                //     .spawn()
                //     .expect("tailmoncli command failed to start");

                // new rust scanner
                // tailmonscan virus --rule /root/projects/tailmon-edr/tailmond/data/rules/current.yara --threads 1 /etc/passwd
                let plugin_root = std::fs::canonicalize("plugins/tailmonscan").unwrap();
                let plugin_root = plugin_root.to_str().unwrap();
                //let pp = format!("{}/tailmonscan", plugin_root);
                let mut child = Command::new(format!("{}/bin/tailmonscan", plugin_root))
                    //.current_dir(plugin_root)
                    .env("LD_LIBRARY_PATH", format!("$LD_LIBRARY_PATH:{}/lib", plugin_root))
                    .arg("virus")
                    .arg("--rule")
                    .arg("data/rules/current.yara")
                    .arg("--threads")
                    .arg("1")
                    .arg("/tmp")
                    .arg("/usr/local")
                    .arg("/root")
                    .arg("/home")
                    .spawn()
                    .expect("tailmoncli command failed to start");

                match rocksdb.find("pid-tailmoncli") {
                    Some(pid) => {
                        let pid = pid.parse::<i32>().unwrap();
                        unsafe {
                            libc::kill(pid, 9);
                        }
                    },
                    None => (),
                }
                rocksdb.save("pid-tailmoncli", child.id().to_string().as_str());

                child.wait().expect("command wasn't running");
                    
                println!("task stop");
                rocksdb.delete("scanning");
            });
        },
    }

    Ok(HttpResponse::Ok().json(models::SysResult {
        success: true,
        error_code: "".to_string(),
        error_message: "".to_string(),
        data: Some(()),
    }))
}



/// Inserts new user with name defined in form.
#[get("/api/monitor/query/info")]
async fn monitor_query_info() -> Result<HttpResponse, Error> {

    let logical_count = heim::cpu::logical_count().await.unwrap();
    let mut frequency: u64 = 1;
    let frequency_result = heim::cpu::frequency().await;
    if frequency_result.is_ok() {
        frequency = frequency_result.unwrap().max().unwrap().get::<heim::units::frequency::gigahertz>();
    }
    
    let platform = heim::host::platform().await.unwrap();
    //let virt = heim_virt::detect().await.unwrap();
    let memory = heim::memory::memory().await.unwrap().total().get::<heim::units::information::byte>();
    let memory_free = heim::memory::memory().await.unwrap().free().get::<heim::units::information::byte>();
    let uptime = heim::host::uptime().await.unwrap().get::<heim::units::time::second>();
    let users = heim::host::users()
        .await.unwrap()
        .try_fold(0usize, |acc, _| async move { Ok(acc + 1) })
        .await.unwrap();

    let partitions = heim::disk::partitions_physical().await.unwrap();
    futures::pin_mut!(partitions);

    let mut disk: u64 = 0;
    let mut disk_free: u64 = 0;
    while let Some(part) = partitions.next().await {
        let part = part.unwrap();
        let usage = part.usage().await.unwrap();
        disk += usage.total().get::<heim::units::information::byte>();
        disk_free += usage.free().get::<heim::units::information::byte>();
    }
    
    Ok(HttpResponse::Ok().json(models::SysResult {
        success: true,
        error_code: "".to_string(),
        error_message: "".to_string(),
        data: Some(models::SystemInfo {
            system: platform.system().to_string(),
            release: platform.release().to_string(),
            version: platform.version().to_string(),
            hostname: platform.hostname().to_string(),
            arch: platform.architecture().as_str().to_string(),
            virt: "".to_string(),
            memory: memory,
            memory_free: memory_free,
            disk: disk,
            disk_free: disk_free,
            uptime: uptime,
            users: users,
            cores: logical_count,
            frequency: frequency,
        })
    }))
}


#[post("/api/login/account")]
async fn login_account(
    local_session: web::Data<session::MemorySession>,
    form: web::Json<models::LoginAccount>,
) -> Result<HttpResponse, Error> {
    let username = match LOCAL_CONFIG.get("username") {
        Some(v) => v.to_string(),
        None => "admin".to_string(),
    };
    let password = match LOCAL_CONFIG.get("password") {
        Some(v) => v.to_string(),
        None => Uuid::new_v4().to_string(),
    };
    if !form.username.eq(username.as_str()) || !form.password.eq(password.as_str()) {
        Ok(HttpResponse::Ok().json(models::LoginResult {
            status: "error".to_string(),
            current_authority: "guest".to_string(),
            login_type: form.login_type.to_string(),
            access_token: "".to_string()
        }))
    } else {
        let access_token = Uuid::new_v4().to_string();
        local_session.session.lock().unwrap().insert(access_token.clone(), Local::now().timestamp_millis());
        Ok(HttpResponse::Ok().json(models::LoginResult {
            status: "ok".to_string(),
            current_authority: "admin".to_string(),
            login_type: form.login_type.to_string(),
            access_token: access_token.clone(),
        }))
    }
}

#[post("/api/login/outLogin")]
async fn login_out_login(
    req: HttpRequest,
    local_session: web::Data<session::MemorySession>,
) -> Result<HttpResponse, Error> {

    let header_option = req.headers().get("Access-Token");
    if header_option.is_some() {
        let header_value_result = header_option.unwrap().to_str();
        if header_value_result.is_ok() {
            let header_value = header_value_result.unwrap();
            let mut data = local_session.session.lock().unwrap();
            data.remove(header_value);
        }
    }

    Ok(HttpResponse::Ok().json(models::SysResult {
        success: true,
        error_code: "".to_string(),
        error_message: "".to_string(),
        data: Some(()),
    }))
}

#[get("/api/currentUser")]
async fn current_user() -> Result<HttpResponse, Error> {
    Ok(HttpResponse::Ok().json(models::SysResult {
        success: true,
        error_code: "".to_string(),
        error_message: "".to_string(),
        data: Some(models::CurrentUser {
            name: "admin".to_string(),
            avatar: "/user.svg".to_string(),
            userid: "00000001".to_string(),
            email: "antdesign@alipay.com".to_string(),
            signature: "海纳百川，有容乃大".to_string(),
            title: "交互专家".to_string(),
            group: "Tailmon-EDR".to_string(),
            tags: vec![
                models::UserTag {
                    key: "0".to_string(),
                    label: "很有想法的".to_string()
                },
                models::UserTag {
                    key: "1".to_string(),
                    label: "专注设计".to_string()
                }
            ],
            notify_count: 12,
            unread_count: 11,
            country: "China".to_string(),
            access: "admin".to_string(),
            address: "西湖区工专路 77 号".to_string(),
            phone: "0752-268888888".to_string(),
            geographic: models::Geographic {
                province: models::LabelKey {
                    label: "浙江省".to_string(),
                    key: "330000".to_string()
                },
                city: models::LabelKey {
                    label: "杭州市".to_string(),
                    key: "330100".to_string()
                }
            }
        })
    }))
}

//use nom::{IResult, bytes::streaming::take};

use nom::{IResult, Err, Needed};

fn take4(i: &[u8]) -> IResult<&[u8], &[u8]>{
  if i.len() < 4 {
    Err(Err::Incomplete(Needed::new(4)))
  } else {
    Ok((&i[4..], &i[0..4]))
  }
}

fn take_len(i: &[u8], len: usize) -> IResult<&[u8], &[u8]>{
    if i.len() < len {
        Err(Err::Incomplete(Needed::new(len)))
    } else {
        Ok((&i[len..], &i[0..len]))
    }
}

#[post("/api/rule/upload")]
async fn upload_rule(
    rocksdb: web::Data<kv::RocksDB>, 
    addr: web::Data<actix::Addr<plugin::suricata_plugin::SuricataPlugin>>, 
    mut payload: Multipart
) -> Result<HttpResponse, Error> {
    // iterate over multipart stream
    while let Some(mut field) = payload.try_next().await? {
        // A multipart/form-data stream has to contain `content_disposition`
        let content_disposition = field.content_disposition();

        let filename = content_disposition
            .get_filename()
            .map_or_else(|| Uuid::new_v4().to_string(), sanitize_filename::sanitize);
        let filepath = format!("./tmp/{}", filename);

        // File::create is blocking operation, use threadpool
        let fp = filepath.clone();
        let mut f = web::block(|| std::fs::File::create(fp)).await??;

        // Field in turn is stream of *Bytes* object
        while let Some(chunk) = field.try_next().await? {
            // filesystem operations are blocking, we have to use threadpool
            f = web::block(move || f.write_all(&chunk).map(|_| f)).await??;
        }

        let content_bytes = std::fs::read(&filepath)?;
        let s = content_bytes.as_slice();
        
        let t = s[0..1][0];
        let rule_filename;
        if t == b'Y' {
            rule_filename = "current.yara";
        } else if t == b'S' {
            rule_filename = "current.suricata";
        } else if t == b'F' {
            rule_filename = "current.falco";
        } else {
            return Ok(HttpResponse::Ok().into());
        }

        let lv = &s[1..];

        let r = take4(lv).unwrap();

        let len: &[u8] = r.1;
        
        let b = i32::from_be_bytes(<[u8; 4]>::try_from(len).expect("Ups, I did it again..."));
        println!("{}", b);
        let rr = take_len(r.0, b.try_into().unwrap()).unwrap();
        // let s1 = String::from_utf8_lossy(rr.0);
        // let s2 = String::from_utf8_lossy(rr.1);
        // println!("result: {}", s1);
        // println!("result: {}", s2);

        let mut file = std::fs::File::create("data/rules/".to_string() + rule_filename)?;
        file.write_all(rr.0)?;

        // let mut file = std::fs::File::create("data/rules/current.meta")?;
        // file.write_all(rr.1)?;
        
        //let contents = std::fs::read_to_string(filepath).expect("{}");
        
        let metadata: model::rule_meta_model::Metadata = serde_json::from_slice(rr.1)?;
        for (rule_id, meta) in metadata.metadata.iter() {
            let rule_json = serde_json::to_string(&meta).unwrap();
            rocksdb.save(&format!("{}-{}", metadata.engine, rule_id), rule_json.as_str());
        }

        let result = addr.send(plugin::suricata_plugin::Reload).await;
        println!("{:?}", result);
        
    }

    Ok(HttpResponse::Ok().into())
}

/// WebSocket handshake and start `MyWebSocket` actor.
async fn echo_ws(req: HttpRequest, stream: web::Payload) -> Result<HttpResponse, Error> {
    ws::start(MyWebSocket::new(), &req, stream)
}


/// WebSocket handshake and start `AlertPromptWebSocket` actor.
// async fn alert_prompt_ws(
//     alert_prompt: web::Data<alert_prompt_ws::AlertPromptWebSocket>, 
//     req: HttpRequest, 
//     stream: web::Payload
// ) -> Result<HttpResponse, Error> {
//     let aaa = *alert_prompt.as_ref();
//     ws::start(aaa, &req, stream)
// }


lazy_static! {

    // set up LocalConfig
    static ref LOCAL_CONFIG: config::LocalConfig = config::LocalConfig::init("config/tailmond.toml");

    
    // static ref SURICATA_CHILD: std::process::Child = Command::new("plugins/suricata/bin/suricata")
    //         .arg("-c")
    //         .arg("plugins/suricata/conf/suricata.yaml")
    //         .arg("-s")
    //         .arg("data/rules/current.suricata")
    //         .arg("-i")
    //         .arg("ens33")
    //         .stdout(Stdio::piped())
    //         .spawn()
    //         .expect("Failed to start suricata process");

    // static ref FALCO_CHILD: std::process::Child = Command::new("plugins/falco/bin/falco")
    //         .arg("-c")
    //         .arg("plugins/falco/conf/falco.yaml")
    //         .arg("-r")
    //         .arg("data/rules/current.falco")
    //         .stdout(Stdio::piped())
    //         .spawn()
    //         .expect("Failed to start suricata process");

    
}



#[actix_web::main]
async fn main() -> std::io::Result<()> {
    dotenv::dotenv().ok();
    env_logger::init_from_env(env_logger::Env::new().default_filter_or("info"));

    // set up database connection pool
    let conn_spec = "data/local.db";
    let manager = ConnectionManager::<SqliteConnection>::new(conn_spec);
    let pool = r2d2::Pool::builder()
        .build(manager)
        .expect("Failed to create pool.");

    log::info!("starting HTTP server at http://localhost:8080");

    

    // set up rocksDB
    let db: kv::RocksDB = kv::KVStore::init("data/kv.db");

    db.delete("scanning");

    // Start Plugins
    //start_plugins();
    match db.find("pid-suricata") {
        Some(pid) => {
            let pid = pid.parse::<i32>().unwrap();
            unsafe {
                libc::kill(pid, 9);
            }
        },
        None => (),
    }
    
    println!("Plugin suricata started!");
    // match db.find("pid-falco") {
    //     Some(pid) => {
    //         let pid = pid.parse::<i32>().unwrap();
    //         unsafe {
    //             libc::kill(pid, 9);
    //         }
    //     },
    //     None => (),
    // }
    // db.save("pid-falco", FALCO_CHILD.id().to_string().as_str());
    // println!("Plugin falco started!");

    // set up LocalConfig
    // let local_config: config::LocalConfig = config::LocalConfig::init("config/tailmond.toml");

    // set up rocksDB
    let ip_block: ip_block::Iptables = ip_block::IpBlock::init("plugins/ipset");

    let bind = match LOCAL_CONFIG.get("bind") {
        Some(addr) => addr.to_string(),
        None => "127.0.0.1".to_string(),
    };

    let port = match LOCAL_CONFIG.get("port") {
        Some(port) => match port.parse::<i32>() {
            Ok(p) => p,
            Err(err) => {
                println!("{:?}", err);
                8080
            },
        },
        None => 8080,
    };

    // set up LocalSession
    let local_session = session::MemorySession::init();

    // Start FalcoScheduler
    // FalcoScheduler::get_scheduler(db.clone(), pool.clone()).start();

    // Start SuricataScheduler
    SuricataScheduler::get_scheduler(db.clone(), pool.clone()).start();

    // Start VirusScheduler
    let last_task = Arc::new(Mutex::new(YaraResult {
        matches: false,
        timestamp: Utc::now().timestamp_millis(),
        mal_obj_type: models::MalObjType::File,
        mal_obj_key: "".to_string(),
        rule_id: Option::None,
        task_total: 0,
        task_finish: 0,
        task_match: 0,
    }));
    VirusScheduler::get_scheduler(db.clone(), pool.clone(), last_task.clone()).start();

    // Start Session Scheduler
    Scheduler::get_scheduler(local_session.clone()).start();

    let suricata_plugin = plugin::suricata_plugin::SuricataPlugin::new();
    db.save("pid-suricata", suricata_plugin.get_pid().to_string().as_str());
    let addr: actix::Addr<plugin::suricata_plugin::SuricataPlugin> = suricata_plugin.start();

    // Send Ping message.
    // send() message returns Future object, that resolves to message result
    // let result = addr.send(Ping).await;

    // Start Web Server
    HttpServer::new(move || {
        App::new()
            // set up DB pool to be used with web::Data<Pool> extractor
            .app_data(web::Data::new(pool.clone()))
            .app_data(web::Data::new(db.clone()))
            //.app_data(web::Data::new(local_config.clone()))
            .app_data(web::Data::new(local_session.clone()))
            .app_data(web::Data::new(ip_block.clone()))
            .app_data(web::Data::new(last_task.clone()))
            .app_data(web::Data::new(addr.clone()))
            // middleware
            .wrap(middleware::Logger::default())
            .wrap(redirect::CheckLogin::set_session(local_session.clone()))
            .wrap(middleware::Compress::default())
            // websocket route
            .service(web::resource("/ws").route(web::get().to(echo_ws)))
            // service route
            .service(user_service::get_user)
            .service(find_alert_by_page)
            .service(login_account)
            .service(current_user)
            .service(login_out_login)
            .service(alert_query_kv)
            .service(monitor_query_info)
            .service(upload_rule)
            .service(solve_alert)
            .service(virus_query_progress)
            .service(virus_start_scan)
            .service(allow_list_service::find_by_page)
            .service(allow_list_service::insert)
            .service(allow_list_service::delete)
            .service(allow_list_service::find_by_id)
            .service(recovery_list_service::find_by_page)
            .service(recovery_list_service::insert)
            .service(recovery_list_service::delete)
            .service(recovery_list_service::find_by_id)
            .service(user_service::notice)
            .service(user_service::activities)
            .service(user_service::fake_workplace_chart_data)
            .service(user_service::notices)
            // static file route
            .service(web::resource("/").route(web::get().to(index)))
            .service(web::resource("/{_:.*}").route(web::get().to(dist)))
            
    })
    .workers(2)
    .bind(bind + ":" + &port.to_string())?
    .run()
    .await
}